package discovery

import (
	"context"
	"net/http"
	"sync"
	"time"
)

type ServiceDiscovery struct {
	services    map[string][]string
	healthCheck bool
	interval    time.Duration
	client      *http.Client
	mu          sync.RWMutex
}

func NewServiceDiscovery(healthCheck bool, interval time.Duration) *ServiceDiscovery {
	return &ServiceDiscovery{
		services:    make(map[string][]string),
		healthCheck: healthCheck,
		interval:    interval,
		client:      &http.Client{Timeout: 5 * time.Second},
	}
}

func (sd *ServiceDiscovery) Register(service string, instances []string) {
	sd.mu.Lock()
	defer sd.mu.Unlock()
	sd.services[service] = instances

	if sd.healthCheck {
		go sd.startHealthCheck(service)
	}
}

func (sd *ServiceDiscovery) GetInstances(service string) []string {
	sd.mu.RLock()
	defer sd.mu.RUnlock()
	return sd.services[service]
}

func (sd *ServiceDiscovery) startHealthCheck(service string) {
	ticker := time.NewTicker(sd.interval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			sd.checkHealth(service)
		}
	}
}

func (sd *ServiceDiscovery) checkHealth(service string) {
	sd.mu.RLock()
	instances := sd.services[service]
	sd.mu.RUnlock()

	var healthyInstances []string
	for _, instance := range instances {
		if sd.isHealthy(instance) {
			healthyInstances = append(healthyInstances, instance)
		}
	}

	sd.mu.Lock()
	sd.services[service] = healthyInstances
	sd.mu.Unlock()
}

func (sd *ServiceDiscovery) isHealthy(instance string) bool {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	req, err := http.NewRequestWithContext(ctx, "GET", instance+"/health", nil)
	if err != nil {
		return false
	}

	resp, err := sd.client.Do(req)
	if err != nil {
		return false
	}
	defer resp.Body.Close()

	return resp.StatusCode == http.StatusOK
}
